#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "itpack2d.h"

float grad(float *data, int columns, int rows, int k,int l)
{
  // A = (sqrt(2)-1)/(2-sqrt(2))
  // B = 1/(2+4*A)
  // C = A*B
  float C = 0.14645, B = 0.20711;
  float x,y, r;
  x = ( -C*data[k*columns+l - columns -1] 
	-B*data[k*columns+l - columns]
	-C*data[k*columns+l - columns +1]
	+C*data[k*columns+l + columns -1] 
	+B*data[k*columns+l + columns]
	+C*data[k*columns+l + columns +1]);
  y = ( -C*data[k*columns+l -1 - columns] 
	-B*data[k*columns+l -1]
	-C*data[k*columns+l -1 + columns]
	+C*data[k*columns+l +1 - columns] 
	+B*data[k*columns+l +1]
	+C*data[k*columns+l +1 + columns]);
  r = sqrt(x*x+y*y);
  return r;
}


float lambda = 5.0;
float tau = 0.5;
float delta = 1.0;

float g(float s)
{
  s = s*s;
  if(s< 0.001) return 1.0;
  else return 1-expf(-3.315/powf(s/lambda,4));
}

#define JCOEF(i,j) (jcoef[(j)*(N)+(i)])
#define COEF(i,j) (coef[(j)*(N)+(i)])
#define NOFFSET(k,l,x,y) (((k)+(x))*(columns)+(l)+(y))

int diffuse(float *data, int columns, int rows)
{
  int t,MAXT = 20;
  int N = columns*rows, NDIM=N, MAXNZ=9;
  int ITMAX = 20, NW=6*N+4*ITMAX, IER;
  float *coef, *WKSP, RPARM[12], *RHS, *G;
  int *jcoef, *IWKSP, IPARM[12];
  int patten[9] = {-columns-1, -columns, -columns+1,
		   -1, 0, 1,
		   +columns-1, +columns, +columns+1,};
  jcoef = malloc(sizeof(int) * N * MAXNZ);
  coef = malloc(sizeof(float) * N * MAXNZ);
  RHS = malloc(sizeof(float) * N);
  G = malloc(sizeof(float) * N); //g(grad(u))
  WKSP = malloc(sizeof(float) * NW);
  IWKSP = malloc(sizeof(int) * N * 3);
  
  dfault_(IPARM, RPARM);
  IPARM[0] = ITMAX;
  IPARM[1] = 1; //output level
  IPARM[4] = 2; // nonsymmetic
  IPARM[11] = 1; //print DIGIT1 and DIGIT2
  
  
  int i,j,k,l;
  for(i=0;i<columns;i++)
    {
      JCOEF(i,0) = i+1;
      JCOEF(N-i-1,0) = N-i;
      COEF(i, 0) =  1.;
      COEF(N-i-1,0) = 1.;
      for(j=1;j<MAXNZ;j++)
	JCOEF(i,j) = 0, JCOEF(N-i-1,j) = 0;
      
    }
  for(i=1;i<rows-1;i++)
    {
      JCOEF(i*columns,0) = i*columns+1;
      COEF(i*columns,0) = 1.;
      JCOEF(i*columns+columns-1,0) = i*columns+columns;
      COEF(i*columns+columns-1,0) = 1.;
      for(j=1;j<MAXNZ;j++)
	JCOEF(i*columns,j) = 0, JCOEF(i*columns+columns-1,j) = 0;
    }


  for(t = 0; t<MAXT; t++)
    {
      memcpy(RHS, data, sizeof(float)*N);
      for(k=1;k<rows-1;k++)
	for(l=1;l<columns-1;l++)
	  {
	    for(j=0;j<MAXNZ;j++)
	      JCOEF(k*columns+l,j) = k*columns + l + patten[j]+1;
	    COEF(k*columns+l,4) = 1.; 
	  }

  
      //cauculate the G 
      for(k=1;k<rows-1;k++)
	for(l=1;l<columns-1;l++)
	  G[k*columns+l] = g(grad(data, columns, rows, k,l));
      for(l=1;l<columns-1;l++)
	{
	  G[l] = G[columns+l]; 
	  G[rows*columns-l-1] = G[(rows-1)*columns-l-1];
	}
       for(k=1;k<rows-1;k++)
	{
	  G[k*columns] = G[k*columns+1]; 
	  G[(k+1)*columns-1] = G[(k+1)*columns-2];
	}
       G[0] = (G[1] + G[columns])/2;
       G[columns-1] = (G[columns-2]+ G[2*columns-1])/2;
       G[(rows-1)*columns] = (G[(rows-2)*columns] + G[(rows-1)*columns+1])/2;
       G[rows*columns-1] = (G[(rows-2)*columns-1] + G[rows*columns-2])/2;

      for(k=1;k<rows-1;k++)
	for(l=1;l<columns-1;l++)
	  {
	    for(j=0;j<MAXNZ;j++)
	      COEF(k*columns+l,j) = 0.;
	    COEF(k*columns+l,4) = 1.;
	    /*
	    for(j=0;j<MAXNZ;j++)
	      COEF(k*columns+l,j) 
		-= g(grad(data, columns, rows, k, l)) * laplace[j]; 
	    */
	    COEF(k*columns+l,1) -= 
	      tau * (G[NOFFSET(k,l,-1,0)] + G[NOFFSET(k,l,0,0)])/2;
	    COEF(k*columns+l,3) -=
	      tau * (G[NOFFSET(k,l,0,-1)] + G[NOFFSET(k,l,0,0)])/2;
	    COEF(k*columns+l,5) -=
	      tau * (G[NOFFSET(k,l,0, 1)] + G[NOFFSET(k,l,0,0)])/2;
	    COEF(k*columns+l,7) -=
	      tau *(G[NOFFSET(k,l,1, 0)] + G[NOFFSET(k,l,0,0)])/2;
	    COEF(k*columns+l,4) -= 
	      COEF(k*columns+l,1)
	      + COEF(k*columns+l,3)
	      + COEF(k*columns+l,5)
	      + COEF(k*columns+l,7);
	  }

      ssorcg_(&N, &NDIM, &MAXNZ, jcoef, coef, RHS, data,
	      IWKSP, &NW, WKSP, IPARM, RPARM, &IER);
    }

  free(jcoef);
  free(coef);
  free(RHS);
  free(WKSP);
  free(IWKSP);
  return 0;
}
